function [x,nrmk,it,nrm,efr] = solverPI_mc(in)
% function [x,normk,it,nrm] = solverPI(input)
%   Solves the coarse mesh diffusion problem via the power 
%   iteration technique.
%
% needs:
%   * function to compute leakage
%   * a way to normalize fission rates
%   * a way to plot outgoing partial currents for a given element
% j. roberts 2/22/2010
% updated 11/1/2011 to use *sparse*  things better

if(in.maxit > 1)
    disp('-------------------------------------------')
    disp('--       beginning power iterations      --')
    disp('-------------------------------------------')
end


x       = in.x0;
it      = 1;
k1      = x(end-1);

% Get response matrices for initial k
[R, F, A, L] = combresp2D(in, x);

% Get initial nonlinear residual norm
nrm(1)  = norm( resid2D(in, R, F, A, L, x) );
nrmk(1) = 1;

%-------------------------------------------------------------------------------
% PERFORM OUTER ITERATIONS
%
%while nrm(it) > in.tol(1)*nrm(it) + in.tol(2) && it < in.maxit+1

while it < in.maxit
    
    fprintf(1,'it = %d, keff = %12.10f, lambda = %12.10f, norm = %12.10e\n', ...
        it, x(end-1), x(end), nrm(it));

    if (nrm(it) < 1e-9)
        break
    end
    %===========================================================================
    % PERFORM INNER ITERATIONS                                                 %
                                                                               %
    MR = in.M*R; % Set a single operator.                                      %
    opts.disp = 0;
    [j1, lambda1] = eigs(MR, 1, 'LM', opts);                                   %
    j1 = j1*sign(sum(j1));                                                     %
    %===========================================================================
%     e = sort(eigs(MR,4),'descend'); rho = abs(e(2)/e(1));
%     fprintf(1,'e1 = %12.10f, e2 = %12.10f, e3 = %12.10f, rho = %12.10f\n', ...
%         e(1), e(2), min(e), rho);
    
    %===========================================================================
    % EIGENVALUE UPDATE                                                    
    %
    k0       = k1;                              % store keff
    j0       = j1;                              % store currents
    lambda0  = lambda1;                         % store current eigenvalue
    gain     = F*j1;                            % compute gains
    absorb   = A*j1;                            % absorption
    leak     = leakage(in, L, [j1;k1;lambda1]); % leakage
    loss     = leak + absorb;                   % and losses
    leak
    k1       = gain / loss;                     % --> keff = gains/losses
    
    % Update unknown vector
    x        = [j1; k1; lambda1];
    
    % Update response matrices for new x
    [R, F, A, L] = combresp2D(in, x); 
    it           = it + 1;
    
    % Update norm
    nrm(it)        = norm( resid2D(in, R, F, A, L, x) );
    normk(it)      = norm(k1-k0);
    
    
    
end

% Fission distribution
fr = F'.*j1;

% degrees of free per element
dofperel = in.numg * (1+in.order) * 4;

efr = zeros(in.numel,1);
for e = 1:in.numel
    for i = 1:dofperel
        idx = (e-1)*dofperel + i;
        efr(e) = efr(e)+fr(idx);
    end
end
efr = efr/norm(efr);

end

function z = resid2D(in, R, F, A, L, x)
% function z = resid2D(x,input)
% This function computes the nonlinear residual
% Here, x(1:end-2) corresponds to the legendre coefficients
%       x(end-1) is keff
%       x(end)   is the current eigenvalue (lambda)
zJ = in.M*(R*x(1:end-2)) - x(end)*x(1:end-2);  % (MR-lambda*I)J = 0
zk = F*x(1:end-2) - x(end-1)*(A*x(1:end-2)+leakage(in, L, x));
zL = 0.5 - 0.5*x(1:end-2)'*x(1:end-2);
z  = [zJ;zk;zL];
end % resid2D


function [R, F, A, L] = combresp2D(in, x)
M     = in.M;
order = in.order;
numel = in.numel;
numg  = in.numg;
% initialize response matrices
Rlist = cell(1,in.numel);
Llist = cell(1,in.numel);
F = zeros( 1, numel*numg*(order+1)*4 );
A = zeros( 1, numel*numg*(order+1)*4 );
lenR = numg*(order+1)*4;
% Get response matrices
if in.locsolv == 1
[Rb Fb Ab Lb]  = mesh_center_2d_resp(x(end-1),in);
else
[Rb Fb Ab Lb]  = resp_kernel(x(end-1),in);
end
% this gets the info for each element type; now, need to place
I = length(in.elements(:,1));
J = length(in.elements(1,:));
k = 0;
for j = 1:J      % rows
    for i = 1:I  % cols
        
        if in.elements(i,j) > 0
            % I am the kth element
            k = k + 1;
            type = in.elements(i,j);
        else
            % else I'm just void
            continue
        end
        tmp(1:lenR,1:lenR)             = Rb(type,:,:); %
        Rlist{k}                       = sparse(tmp);  %
        F(1, ((k-1)*lenR+1):(k*lenR) ) = Fb(type,:);   %
        A(1, ((k-1)*lenR+1):(k*lenR) ) = Ab(type,:);   %
        myL(:,:) = Lb(type,:,:);
        my1 = myL(1,:);
        my2 = myL(2,:);
        my3 = myL(3,:);
        my4 = myL(4,:);
        % for an element, we have
        Llist{k}  = sparse(...
            [my1 my2 my3 my4
             my2 my1 my4 my3
             my4 my3 my1 my2
             my3 my4 my2 my1]);
    end
end
% Build R and the leakage operator
R   = blkdiag(Rlist{:});
L   = blkdiag(Llist{:});
end % combresp2D

function l = leakage(in, L, x)
% build mm
lenM = length(in.M(:,1));
mm = zeros(lenM, 1 );
for i = 1:lenM
    if ( sum(in.M(:,i)) + sum(in.M(i,:)) )==0
        mm(i)=1;
    else
        mm(i)=0;
    end
end
mm_v     = mm(1:in.numg*(in.order+1):end);
mm_v_len = length(mm_v);
mm       = spdiags(mm_v, 0, mm_v_len, mm_v_len);
l = sum( mm*(L*x(1:end-2)) );
end
